#ifndef OPERATIONAL_SPACE_CONTROLLER_H
#define OPERATIONAL_SPACE_CONTROLLER_H

#include "RobotQuantity.h"

//task is x'' = ddxdes
struct AccelerationTask
{
  SmartPointer<RobotQuantity> q;
  Vector ddxdes;
  Real weight;
};

//task is T[active] = Tdes
struct TorqueTask
{
  vector<int> active;
  Vector Tdes;
  Real weight;
};

//Task is Af = fdes
struct ForceTask
{
  Matrix A;
  Vector fdes;
  Real weight;
};


struct OperationalSpaceController : public ControllerBase
{
  vector<AccelTask> accelTasks;
  vector<TorqueTask> torqueTasks;
  vector<ForceTask> forceTasks;


  OperationalSpaceController(Robot& _robot,RobotSensors& _sensors,RobotCommands& _commands)
    :ControllerBase(_robot,_sensors,_commands)
  {}

  virtual void UpdateTasks()=0;
  virtual void EstimateStance(Stance& s)=0;

  virtual void Cycle(Real dt) {
    UpdateRobot();
    UpdateTasks();
    Stance s;
    EstimateStance(s);
    Vector torques;
    TasksToTorques(s,torques);
    CommandTorques(torques);
  }

  void TasksToTorques(const Stance& estimatedStance,Vector& t) {
    t.resize(robot.kNumDof);
    for(int i=0;i<6;i++) t(i)=0;
    Vector t_active;
    t_active.setRef(t,6,1,robot.kNumActuators);

    //use torques that closely satisfy ddq
    Matrix B;
    Assert(robot.JP.m == robot.q.n);
    Assert(robot.JP.n == robot.q.n);
    robot.GetKineticEnergyMatrix(B);

    /*
      q'' = B^-1(ST + Jtf - CG)
      solve min{f,T} ||q''-q''des||
      s.t. |T|<=Tmax
           Jq'' = 0   (ensures that contact point accelerations are 0)
           fi in FCi

      expanding...
      q''-q''des = B^-1(ST + Jtf) - (q''des + B^-1CG)
      Jq'' = JB^-1(ST+Jtf) - JB^-1 CG
    */
    if(curStance.empty()) {
      B.mul(ddq,t);
      return;
    }

    CholeskyDecomposition<Real> LL;
    Matrix Binv;
    bool r=LL.set(B);
    Assert(r==true);
    LL.getInverse(Binv);

    Vector CG,BinvCG;
    robot.GetGravityTorques(kGravityVector,CG);
    Binv.mul(CG,BinvCG);
    //TODO coriolis force estimate?
    int n=robot.q.n;
    int nc=NumContactPoints(curStance);
    int nt=robot.kNumActuators;
    Matrix J(nc*3,robot.q.n);
    int m=0;
    for(Stance::const_iterator k=curStance.begin();k!=curStance.end();k++) {
      const Hold& h=k->second;
      for(size_t i=0;i<h.contacts.size();i++,m+=3) {
	Matrix Jm;
	Jm.setRef(J,m,0,1,1,3,robot.q.n);
	Vector3 localPos;
	robot.links[h.link].T_World.mulInverse(h.contacts[i].x,localPos);
	robot.GetPositionJacobian(localPos,h.link,Jm);
      }
    }
    m=0;
    for(Stance::const_iterator k=curStance.begin();k!=curStance.end();k++) {
      const Hold& h=k->second;
      if(h.contacts.size() >= 3) m+=6;
      else if(h.contacts.size() == 2) m+=5;
      else m+=3;
    }
    cout<<"K size: "<<m<<endl;
    Matrix K(m,robot.q.n);
    m=0;
    for(Stance::const_iterator k=curStance.begin();k!=curStance.end();k++) {
      const Hold& h=k->second;
      Matrix Km;
      if(h.contacts.size() >= 3) {
	Km.setRef(K,m,0,1,1,6,robot.q.n);
	robot.GetFullJacobian(Vector3(Zero),h.link,Km);
	m += 6;
      }
      else if(h.contacts.size() == 2) {
	//constrain position and 2 dofs of orientation
	Vector3 axis,x,y;
	axis = h.contacts[1].x - h.contacts[0].x;
	axis.inplaceNormalize();
	axis.getOrthogonalBasis(x,y);
	Vector3 localPos;
	Km.setRef(K,m,0,1,1,3,robot.q.n);
	robot.links[h.link].T_World.mulInverse(h.contacts[0].x,localPos);
	robot.GetPositionJacobian(localPos,h.link,Km);
	Matrix temp(3,robot.q.n);
	robot.links[h.link].T_World.mulInverse(h.contacts[1].x,localPos);
	robot.GetPositionJacobian(localPos,h.link,temp);
	Km.setRef(K,m+3,0,1,1,2,robot.q.n);
	for(int j=0;j<robot.q.n;j++) {
	  Km(0,j) = x.dot(Vector3(temp(0,j),temp(1,j),temp(2,j)));
	  Km(1,j) = y.dot(Vector3(temp(0,j),temp(1,j),temp(2,j)));
	}
	m += 5;
      }
      else {
	Km.setRef(K,m,0,1,1,3,robot.q.n);
	Vector3 localPos;
	robot.links[h.link].T_World.mulInverse(h.contacts[0].x,localPos);
	robot.GetPositionJacobian(localPos,h.link,Km);
	m += 3;
      }
    }
    cout<<"K: "<<endl<<MatrixPrinter(K,MatrixPrinter::AsciiShade)<<endl;

    MinNormProblem lp;
    lp.norm = One;
    lp.A.resize(n,nt+nc*3);
    lp.b.add(BinvCG,ddq);
    Matrix mtemp;
    Vector vtemp;
    //norm: ST+Jtf
    lp.A.setZero();
    mtemp.setRef(lp.A,0,0,1,1,n,nt);
      mtemp.setZero();
      for(int i=0;i<nt;i++)
	mtemp(i+6,i) = One;
    mtemp.setRef(lp.A,0,nt,1,1,n,nc*3);
      mtemp.setTranspose(J);
    //transform by Binv
    mtemp.clear();
    mtemp=lp.A;
    lp.A.mul(Binv,mtemp);
    mtemp.clear();
    //FC
    lp.Aineq.resize(kNumFCEdges*nc,nt+nc*3,Zero);
    lp.bineq.resize(kNumFCEdges*nc,Zero);
    mtemp.setRef(lp.Aineq,0,nt,1,1,nc*kNumFCEdges,nc*3);
    GetFrictionConePlanes(curStance,kNumFCEdges,mtemp);
    //torque bounds
    lp.u.resize(nt+nc*3,Inf);
    lp.l.resize(nt+nc*3,-Inf);
    for(int i=0;i<nt;i++) {
      lp.l(i) = -robot.torqueMax(i);
      lp.u(i) = robot.torqueMax(i);
    }
    //equality constraint
    lp.Aeq.mul(K,lp.A);
    K.mul(BinvCG,lp.beq);

    //weights on the norm?
    for(int i=0;i<lp.A.m;i++) {
      Real scale;
      if(i<3) scale=1;
      else if(i<6) scale=1;
      else scale=1;
      Vector ai;
      lp.A.getRowRef(i,ai);
      ai *= scale;
      lp.b(i) *= scale;
    }

    Assert(lp.IsValid());
    lp.Assemble();
    //lp.Print(cout);
    //cout<<"Linear program: "<<endl;
    //lp.lp.Print();
    for(int i=0;i<6;i++) t(i)=0;
    Vector x;
    LinearProgram::Result res=lp.Solve(x);
    switch(res) {
    case LinearProgram::Feasible:
      x.getSubVectorCopy(0,t_active);
      vtemp.setRef(x,nt,1,nc*3);
      //T = tdes - Jtf
      cout<<"Commanded torques: "<<VectorPrinter(t_active,VectorPrinter::AsciiShade)<<endl;
      cout<<"F: "<<VectorPrinter(vtemp)<<endl;
      cout<<"L"<<lp.norm<<" error: "<<lp.Norm(x)<<endl;
      {
	Vector temp,temp2;
	B.mul(ddq,temp);
	cout<<"original B*ddq: "<<VectorPrinter(temp,VectorPrinter::AsciiShade)<<endl;
	cout<<"original ddq: "<<VectorPrinter(ddq,VectorPrinter::AsciiShade)<<endl;
	J.mulTranspose(vtemp,temp);
	temp += t;
	robot.GetGravityTorques(kGravityVector,temp2);
	temp -= temp2;
	cout<<"New B*ddq: "<<VectorPrinter(temp,VectorPrinter::AsciiShade)<<endl;
	temp2 = temp;
	LL.backSub(temp2,temp);
	cout<<"New ddq: "<<VectorPrinter(temp,VectorPrinter::AsciiShade)<<endl;
	temp -= ddq;
	cout<<"Difference: "<<VectorPrinter(temp,VectorPrinter::AsciiShade)<<endl;
      }
      break;
    default:
      cout<<"Error computing torques! result "<<res<<endl;
      getchar();
      B.mul(ddq,t);
      break;
    }

    /*
    for(int i=0;i<t_active.n;i++) {
      int link=ACTUATOR_TO_LINK(i);
      if(curSpace->freeDof[link]) {
	t_active(i) += kP(i)*DofSub(robot,q_des,robot.q,link);
	t_active(i) += kD(i)*(dq_des(link)-robot.dq(link));
      }
      else {
	t_active(i) += kP(i)*DofSub(robot,q_des,robot.q,link)*10;
	t_active(i) += kD(i)*(dq_des(link)-robot.dq(link));
      }
    }
    */
  }
};
